import time
from Enet2 import Enet, UData,YData
from datetime import datetime
import sys


# 导入自定义函数和CANFD SDK动态链接库>>>
from os.path import dirname, join
import sys
sys.path.append(join(dirname(__file__), "src"))
sys.path.append(dirname(__file__))
from src import *


from imu import IMU
import time
imu=IMU()
time.sleep(2)




# Initialize Enet
enet = Enet()
enet.ReceiveY()
now=datetime.now()




loopPeriod=control_period
cnt=0

time.sleep(1)


site_force = np.zeros(6)
filtered_site_force = np.zeros(6)
filtered_imu_angle = np.zeros(3)


l = 0.3 # 头的宽度
roT_h = np.array([[0, 1, 0, 0],[0, 0, 1, 0],[1, 0, 0, -l],[0, 0, 0, 1]]) #头中心点到右臂o点的变换阵
loT_h = np.array([[0, 1, 0, 0],[0, 0, -1, 0],[-1, 0, 0, -l],[0, 0, 0, 1]]) #头中心点到左臂o点的变换阵
hT_e = np.eye(4)
box_pos = np.array([0, 0.5, -0.4])
box_quat = np.array([1,0,0,0])
imu.angle = [0,0,0]

try:
    
    loopStartTime=time.time()
    while True:
        
        cnt+=1
        start_time = time.time()
        t=start_time-loopStartTime

        # 2.1 ReadY from YPool
        start_time = time.time()

        y_data = enet.ReadY()
        
        q_real_degree = np.array(y_data.pos)
        q_real_degree_l = q_real_degree[0:6]
        q_real_degree_r = q_real_degree[6:12]
        q_real_int_l = trans_q_degree2intq(q_real_degree_l)
        q_real_int_r = trans_q_degree2intq(q_real_degree_r)

        q_real_float_l = trans_intq2q(q_real_int_l)
        q_real_float_r = trans_intq2q(q_real_int_r)

        


        

        
        if t<15:
            box_width = 1
        else:
            box_width = 0.5


        if len(imu.angle) == 3:
            imu_angle_raw = np.array([imu.angle[0],imu.angle[1],imu.angle[2]])
        filtered_imu_angle = alpha * imu_angle_raw + (1 - alpha) * filtered_imu_angle
        anglebias = np.pi*np.clip(filtered_imu_angle.copy()[0], -1, 1)/180

        head_pos = np.array([0,0,0])
        c, s = np.cos(anglebias), np.sin(anglebias)
        head_rot = np.array([
            [1,  0,   0],
            [0,  c,  -s],
            [0,  s,   c]
        ])
        hT_e[:3, 3] = -head_rot.T@head_pos
        hT_e[:3, :3] = head_rot.T


        rd_pos, rd_quat, ld_pos, ld_quat = pick_box_r_l(box_width, box_pos, box_quat)

        roT_e = roT_h@hT_e
        roP_D, roR_D = eD2oD(rd_pos, rd_quat, roT_e)

        loT_e = loT_h@hT_e
        loP_D, loR_D = eD2oD(ld_pos, ld_quat, loT_e)

        # TODO 现在是12维，还得改
        # site_force_raw = np.array(y_data.F)
        
        # filtered_site_force = alpha * site_force_raw + (1 - alpha) * filtered_site_force
        # site_force = filtered_site_force.copy()
        
        
        
        ctrl_l = motor_angle_clip(ik_general(q_real_float_l, loP_D, loR_D))
        ctrl_r = motor_angle_clip(ik_general(q_real_float_r, roP_D, roR_D))

        # ctrl_l = np.zeros(6)
        # ctrl_r = np.zeros(6)

        intikq_l = trans_q2intq(ctrl_l)
        intikq_r = trans_q2intq(ctrl_r)


        ctrl_intq_l = get_interpoint_l(q_real_int_l, intikq_l, max_dq)
        ctrl_intq_r = get_interpoint_r(q_real_int_r, intikq_r, max_dq)

        ctrl_intq = np.hstack((ctrl_intq_l, ctrl_intq_r))
        ctrl_q_degree = trans_intq2q_degree(ctrl_intq).tolist()
        # 添加条件限制打印频率
        if cnt % 100 == 0:  # 每100次循环打印一次
            print(f"Average time cost: {(time.time()-loopStartTime)/cnt:.4f}s")
            print(f"q_real is {q_real_degree}")
            print(f"ctrl_q_degree is {ctrl_q_degree}")
            print(f"ctrl l {ctrl_l} and ctrl_r {ctrl_r}")

        # 2.2 TODO:Compute U from Y
        ts = int(time.time() * 1_000_000)  # current timestamp from epoch (microsecond)
        u_data = UData(timestamp=ts, pos=ctrl_q_degree)

        # 2.3 Send U to remoteIP
        enet.SendU(u_data)

        # Sleep to maintain 1kHz loop
        elapsed_time = time.time() - start_time
        # print(elapsed_time)
        time.sleep(max(0, loopPeriod - elapsed_time))
except KeyboardInterrupt:
    print("Exiting...")